import tensorflow
from tensorflow.python.keras.callbacks import EarlyStopping

import kashgari
from kashgari.tasks.labeling import BiGRU_CRF_Model
from kashgari.embeddings import BERTEmbedding


def read_file(path):
    features = [[]]
    labels = [[]]
    with open(path, 'r', encoding='utf-8') as f:
        message = f.readlines()
    for i in message:
        i = i.split("\t")
        features.append(i[0].split())
        labels.append(i[1].split())
    return features, labels


train_data, train_labels = read_file("medicinal_data_indication/train_data.txt")
test_data, test_labels = read_file("medicinal_data_indication/test_data.txt")
valid_data, valid_labels = read_file("medicinal_data_indication/valid_data.txt")

bert_embed = BERTEmbedding('wwm',
                           trainable=True,
                           task=kashgari.LABELING,
                           sequence_length=10)
model = BiGRU_CRF_Model()
# model = BiGRU_CRF_Model()
early_stopping = EarlyStopping(
    monitor='val_accuracy',
    min_delta=0,
    patience=8,
    verbose=1,
    mode='auto'
)
tb_cb = tensorflow.python.keras.callbacks.TensorBoard(log_dir="log_dir", write_images=1, histogram_freq=1)

model.fit(train_data,
          train_labels,
          valid_data,
          valid_labels,
          epochs=500,
          callbacks=[early_stopping, tb_cb]
          )
model.evaluate(test_data, test_labels)
